# -*- coding: utf-8 -*-
'''
File:configuration_barranzi2.py                                         
Time:2025/9/29 15:43                                    
IDE:PyCharm                                     
Author:Barranzi An                                        
email:awc19930818@outlook.com                           
github:https://github.com/La0bALanG                     
Barranzi's Blog:私聊巴郎子要文档链接                  
requirement:(Please describle your requirement here) -->
'''

# 导入必要的类型注解和工具库
from typing import List, Optional, Callable  # 导入类型提示相关的类，用于静态类型检查
from dataclasses import dataclass, field  # 导入数据类装饰器和字段工具，用于简化类的定义
import torch  # 导入PyTorch深度学习框架，用于张量操作和神经网络构建


@dataclass(kw_only=True)  # 使用数据类装饰器，kw_only=True表示所有字段都必须通过关键字参数传递
class RoPEConfig(object):
    """
    RoPE (Rotary Position Embedding) 配置类
    RoPE是一种相对位置编码方法，通过旋转矩阵来编码位置信息，能够有效处理长序列依赖关系

    Args:
        rope_type (`str`, default is default):
            RoPE的类型，支持三种模式：
            - 'default': 标准RoPE实现
            - 'dynamic': 动态RoPE，能够根据序列长度动态调整位置编码
            - 'yarn': YaRN (Yet another RoPE extensioN) 扩展RoPE，支持更长的上下文长度
        rope_theta (`float`, default is 10000.0):
            RoPE的基础旋转角度参数，控制位置编码的频率
            较大的值会使位置编码变化更缓慢，适合处理长序列
        factor (`float`):
            缩放因子，用于扩展RoPE的上下文长度
            除了'default'类型外的所有RoPE类型都使用此参数
            在大多数缩放类型中，factor为x时，模型能够处理的序列长度为原始最大预训练长度的x倍
        partial_rotary_factor (`float`):
            部分旋转因子，控制应用RoPE的维度比例
            1.0表示所有维度都应用旋转，小于1.0表示部分维度应用旋转
        attention_factor (`float`):
            注意力计算的缩放因子，仅用于'yarn'类型
            如果未指定，将根据factor字段推断建议值并使用实现推荐的默认值
        beta_fast (`float`):
            快速边界参数，仅用于'yarn'类型
            用于在线性斜坡函数中设置外推（仅）的边界
            如果未指定，默认为32
        beta_slow (`float`):
            慢速边界参数，仅用于'yarn'类型
            用于在线性斜坡函数中设置插值（仅）的边界
            如果未指定，默认为1
        mscale (`float`):
            多维度缩放参数，仅用于'yarn'类型
            用于调整不同维度的缩放比例
            默认为None，表示使用自动计算值
        mscale_all_dim (`float`):
            全维度缩放参数，仅用于'yarn'类型
            对所有维度应用统一的缩放
            默认为None，表示使用自动计算值
    """
    # RoPE类型配置，默认为标准实现
    rope_type: str = 'default'
    # RoPE基础旋转角度，默认为10000.0
    rope_theta: float = 10000.0
    # 上下文长度扩展因子，默认为1.0（不扩展）
    factor: float = 1.0
    # 部分旋转因子，默认为1.0（全维度旋转）
    partial_rotary_factor: float = 1.0
    # YaRN快速边界参数，默认为32
    beta_fast: float = 32
    # YaRN慢速边界参数，默认为1
    beta_slow: float = 1
    # 多维度缩放参数，可选，默认为None
    mscale: Optional[float] = None
    # 全维度缩放参数，可选，默认为None
    mscale_all_dim: Optional[float] = None
    # 注意力缩放因子，可选，默认为None
    attention_factor: Optional[float] = None


@dataclass(kw_only=True)  # 使用数据类装饰器，所有字段必须通过关键字传递
class MoEConfig(object):
    """
    MoE (Mixture of Experts) 混合专家模型配置类
    MoE是一种模型架构，通过在推理时动态选择不同的专家网络来处理不同的输入，从而在不显著增加计算成本的情况下扩展模型容量

    Args:
        intermediate_size (`int`, *optional*, defaults to None):
            MoE专家网络的中间层维度大小
            决定了每个专家网络的容量和复杂度
            如果为None，则使用基础模型的中间层大小
        n_dense_layer (`int`, default to None):
            使用密集层而非MoE层的层数
            指定在模型的前几层或后几层使用标准的Transformer密集层
            如果为None，则所有层都使用MoE架构
        num_experts_per_tok (`int`, *optional*, defaults to None):
            每个token选择的专家数量
            控制在推理时为每个输入token激活多少个专家网络
            如果为None，表示使用密集模型（不使用MoE）
        n_shared_experts (`int`, *optional*, defaults to None):
            共享专家的数量
            共享专家会被所有token使用，提供通用的知识表示
            如果为None，表示没有共享专家
        n_routed_experts (`int`, *optional*, defaults to None):
            路由专家的数量
            路由专家通过路由网络被动态选择，处理特定的输入模式
            如果为None，表示没有路由专家（即使用密集模型）
        routed_scaling_factor (`float`, *optional*, defaults to 1.0):
            路由专家的缩放因子
            用于调整路由专家输出的权重，平衡不同专家的贡献
        seq_aux (`bool`, *optional*, defaults to True):
            是否为每个单独的样本计算辅助损失
            辅助损失用于训练路由网络，确保专家负载均衡
            如果为True，则对每个样本单独计算辅助损失
        norm_topk_prob (`bool`, *optional*, defaults to False):
            是否对路由专家的权重进行归一化
            如果为True，则对选中的top-k专家权重进行归一化，确保权重和为1
    """
    # MoE专家网络的中间层维度，可选，默认为None
    intermediate_size: Optional[int] = None
    # 使用密集层的数量，可选，默认为None
    n_dense_layer: Optional[int] = None
    # 每个token选择的专家数量，可选，默认为None
    num_experts_per_tok: Optional[int] = None
    # 共享专家数量，可选，默认为None
    n_shared_experts: Optional[int] = None
    # 路由专家数量，可选，默认为None
    n_routed_experts: Optional[int] = None
    # 路由专家缩放因子，默认为1.0
    routed_scaling_factor = 1.0
    # 是否为每个样本计算辅助损失，默认为True
    seq_aux: bool = True
    # 是否对top-k专家权重进行归一化，默认为False
    norm_topk_prob: bool = False


@dataclass(kw_only=True)  # 使用数据类装饰器，所有字段必须通过关键字传递
class Config(object):
    """
    大型语言模型(LLM)的基础配置类
    包含了Transformer模型架构的核心参数，用于定义模型的结构和训练行为

    Args:
        vocab_size (`int`, *optional*, default is 21128):
            词汇表大小，定义模型能够识别和生成的token数量
            对于中文模型通常包含汉字、标点符号、特殊标记等
        hidden_size (`int`, *optional*, default is 4096):
            隐藏层维度大小，也称为d_model
            决定了模型每层的表示能力和参数数量
            较大的hidden_size能提供更强的表达能力，但会增加计算成本
        intermediate_size (`int`, *optional*, default is 11008):
            前馈神经网络的中间层维度
            通常为hidden_size的2-4倍，用于增加模型的非线性表达能力
        num_hidden_layers (`int`, *optional*, default is 32):
            Transformer解码器的层数
            更多的层数能提供更强的抽象能力，但会增加训练难度和推理时间
        num_attention_heads (`int`, *optional*, default is 32):
            注意力头的数量
            多头注意力机制允许模型同时关注不同位置和不同表示子空间的信息
        num_key_value_heads (`int`, *optional*, default is num_attention_heads):
            键值头的数量，用于分组查询注意力(GQA)机制
            如果num_key_value_heads == num_attention_heads，使用标准多头注意力(MHA)
            如果num_key_value_heads == 1，使用多查询注意力(MQA)
            如果1 < num_key_value_heads < num_attention_heads，使用分组查询注意力(GQA)
            GQA在保持性能的同时能显著减少推理时的内存占用
        max_position_embeddings (`int`, *optional*, default is 2048):
            最大位置嵌入长度，定义模型能处理的最大序列长度
            超过此长度的序列会被截断或需要特殊处理
        original_max_position_embeddings (`int`, *optional*, default is None):
            原始最大位置嵌入长度，用于位置编码扩展
            如果使用了RoPE扩展，此参数保存原始预训练时的最大长度
        attention_dropout (`float`, *optional*, default is 0.1):
            注意力机制的dropout率
            在训练过程中随机丢弃一部分注意力权重，防止过拟合
        attention_implementation (`str`, default is auto):
            注意力机制的实现方式，支持三种模式：
            - 'auto': 自动选择，优先使用F.scaled_dot_product_attention
            - 'sdpa': 强制使用PyTorch的scaled_dot_product_attention实现
            - 'default': 使用纯Python实现的注意力机制
        rope_config (`RoPEConfig`):
            RoPE(旋转位置编码)的配置对象
            包含位置编码相关的所有参数设置
        moe_config (`MoEConfig`):
            MoE(混合专家)的配置对象
            如果为None，表示不使用MoE架构
        use_qk_norm (`bool`):
            是否对查询(Q)和键(K)进行层归一化
            归一化可以稳定训练过程，提高模型性能
    """
    # 词汇表大小，必须指定
    vocab_size: int
    # 隐藏层维度大小，必须指定
    hidden_size: int
    # 前馈网络中间层维度，必须指定
    intermediate_size: int
    # Transformer解码器层数，必须指定
    num_hidden_layers: int
    # 注意力头数量，必须指定
    num_attention_heads: int
    # 键值头数量，必须指定
    num_key_value_heads: int
    # 最大位置嵌入长度，必须指定
    max_position_embeddings: int
    # 原始最大位置嵌入长度，可选，默认为None
    original_max_position_embeddings: Optional[int] = None
    # 注意力dropout率，默认为0.1
    attention_dropout: float = 0.1
    # 注意力实现方式，默认为'auto'
    attention_implementation: str = 'auto'
    # RoPE配置，默认使用默认配置
    rope_config: RoPEConfig = field(default_factory=RoPEConfig)
    # MoE配置，可选，默认为None
    moe_config: Optional[MoEConfig] = None
    # 是否使用QK归一化，默认为True
    use_qk_norm: bool = True

    def __post_init__(self):
        """
        数据类初始化后的后处理方法
        在所有字段初始化完成后自动调用，用于进行额外的验证和设置
        """
        # 如果没有指定键值头数量，则设置为与注意力头数量相同
        # 这确保了默认情况下使用标准的多头注意力(MHA)
        if self.num_key_value_heads is None:
            self.num_key_value_heads = self.num_attention_heads


@dataclass(kw_only=True)  # 使用数据类装饰器，所有字段必须通过关键字传递
class VLMConfig(Config):
    """
    VLM (Vision-Language Model) 视觉语言模型配置类
    继承自基础Config类，添加了视觉处理相关的配置参数
    用于构建能够同时处理图像和文本的多模态模型

    Args:
        image_tok (`int`):
            图像token的ID
            用于在词汇表中标识图像的特殊标记
        image_size (`int`):
            输入图像的尺寸
            通常是正方形图像的边长，如224、384等
        patch_size (`int`):
            图像分块的大小
            Vision Transformer将图像分割成小块进行处理，此参数定义小块的尺寸
        tokens_per_image (`int`):
            每张图像对应的token数量
            等于(image_size // patch_size) ** 2，表示图像被分割成的总块数
        vision_hidden_size (`int`):
            视觉编码器的隐藏层维度
            可能与文本编码器的hidden_size不同，用于处理视觉特征
        vision_tower (`Callable[[torch.Tensor], torch.Tensor]`):
            视觉编码塔(视觉编码器)
            一个可调用对象，接收图像张量作为输入，输出视觉特征表示
            通常是预训练的Vision Transformer或CNN模型
    """
    # 图像token的ID，必须指定
    image_tok: int
    # 输入图像的尺寸，必须指定
    image_size: int
    # 图像分块大小，必须指定
    patch_size: int
    # 每张图像的token数量，必须指定
    tokens_per_image: int
    # 视觉编码器隐藏层维度，必须指定
    vision_hidden_size: int
    # 视觉编码塔，必须指定
    vision_tower: Callable[[torch.Tensor], torch.Tensor]